package tunnel

import (
	"bufio"
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"net"
	"os"
	"sync"
	"time"

	"natproxy/pkg/common"

	"compress/gzip"
)

// Tunnel 表示一个隧道连接
type Tunnel struct {
	conn        net.Conn
	connMux     sync.RWMutex
	done        chan struct{}
	doneMux     sync.Once
	heartbeat   *time.Ticker
	isConnected bool
	reader      *bufio.Reader
}

// FileTransfer 文件传输相关的常量和结构
const (
	maxMessageSize = 100 * 1024 * 1024  // 100MB 最大消息大小
	bufferSize     = 1 * 1024 * 1024    // 1MB 缓冲区大小
	maxFileSize    = 1024 * 1024 * 1024 // 1GB 最大文件大小
	readBufferSize = 32 * 1024          // 32KB
)

// FileTransferStatus 表示文件传输状态
type FileTransferStatus struct {
	FileID       string
	TotalChunks  int
	CurrentChunk int
	TotalSize    int64
	Completed    bool
	Error        error
}

// NewTunnel 创建新的隧道连接
func NewTunnel(conn net.Conn) *Tunnel {
	if tcpConn, ok := conn.(*net.TCPConn); ok {
		tcpConn.SetKeepAlive(true)
		tcpConn.SetKeepAlivePeriod(30 * time.Second)
	}

	return &Tunnel{
		conn:        conn,
		done:        make(chan struct{}),
		isConnected: true,
		reader:      bufio.NewReaderSize(conn, bufferSize),
	}
}

// SendMessage 发送消息
func (t *Tunnel) SendMessage(msg *common.TunnelMessage) error {
	t.connMux.RLock()
	defer t.connMux.RUnlock()

	if t.conn == nil {
		return fmt.Errorf("connection is closed")
	}

	// 压缩消息体
	if err := msg.CompressPayload(); err != nil {
		return fmt.Errorf("compress message failed: %v", err)
	}

	// 序列化消息
	data, err := json.Marshal(msg)
	if err != nil {
		return fmt.Errorf("marshal message failed: %v", err)
	}

	// 检查消息大小
	if len(data) > maxMessageSize {
		return fmt.Errorf("message too large: %d bytes", len(data))
	}

	// 检查数据中是否包含空字符
	if bytes.IndexByte(data, 0) != -1 {
		return fmt.Errorf("message contains null bytes")
	}

	// 添加换行符作为消息分隔符
	data = append(data, '\n')

	// 写入数据
	_, err = t.conn.Write(data)
	if err != nil {
		t.conn = nil
		return fmt.Errorf("write message failed: %v", err)
	}

	return nil
}

// ReadMessage 读取消息
func (t *Tunnel) ReadMessage() (*common.TunnelMessage, error) {
	t.connMux.RLock()
	defer t.connMux.RUnlock()

	if t.conn == nil {
		return nil, fmt.Errorf("connection is closed")
	}

	data, err := t.reader.ReadSlice('\n')
	if err != nil {
		if err != io.EOF {
			t.conn = nil
		}
		return nil, fmt.Errorf("read message failed: %v", err)
	}

	data = bytes.TrimSpace(data)

	// 尝试解压数据
	if len(data) > 2 && data[0] == 0x1f && data[1] == 0x8b {
		reader, err := gzip.NewReader(bytes.NewReader(data))
		if err != nil {
			return nil, fmt.Errorf("create gzip reader failed: %v", err)
		}
		defer reader.Close()

		uncompressed, err := io.ReadAll(reader)
		if err != nil {
			return nil, fmt.Errorf("decompress data failed: %v", err)
		}
		data = uncompressed
	}

	var msg common.TunnelMessage
	if err := json.Unmarshal(data, &msg); err != nil {
		return nil, fmt.Errorf("unmarshal message failed: %v (data: %s)", err, string(data))
	}

	// 解压 payload
	if len(msg.Payload) > 0 {
		reader := bytes.NewReader([]byte(msg.Payload))
		gzReader, err := gzip.NewReader(reader)
		if err == nil {
			defer gzReader.Close()
			if uncompressed, err := io.ReadAll(gzReader); err == nil {
				msg.Payload = uncompressed
			}
		}
	}

	return &msg, nil
}

// StartHeartbeat 开始心跳
func (t *Tunnel) StartHeartbeat(interval time.Duration) {
	t.connMux.Lock()
	if t.heartbeat != nil {
		t.heartbeat.Stop()
	}
	t.heartbeat = time.NewTicker(interval)
	t.connMux.Unlock()

	go func() {
		for {
			select {
			case <-t.done:
				return
			case <-t.heartbeat.C:
				if err := t.SendMessage(&common.TunnelMessage{Type: common.TypeHeartbeat}); err != nil {
					t.Close()
					return
				}
			}
		}
	}()
}

// Close 关闭隧道
func (t *Tunnel) Close() error {
	t.connMux.Lock()
	defer t.connMux.Unlock()

	t.doneMux.Do(func() {
		close(t.done)
	})

	if t.heartbeat != nil {
		t.heartbeat.Stop()
		t.heartbeat = nil
	}

	if t.conn != nil {
		err := t.conn.Close()
		t.conn = nil
		t.isConnected = false
		return err
	}

	return nil
}

// IsConnected 检查隧道是否连接
func (t *Tunnel) IsConnected() bool {
	if t == nil {
		return false
	}
	t.connMux.RLock()
	defer t.connMux.RUnlock()
	return t.conn != nil && t.isConnected
}

// ReceiveFile 接收文件
func (t *Tunnel) ReceiveFile(savePath string, progressChan chan<- FileTransferStatus) error {
	file, err := os.Create(savePath)
	if err != nil {
		return fmt.Errorf("create file failed: %v", err)
	}
	defer file.Close()

	var status FileTransferStatus
	chunks := make(map[int][]byte)

	for {
		msg, err := t.ReadMessage()
		if err != nil {
			return fmt.Errorf("read message failed: %v", err)
		}

		switch msg.Type {
		case common.TypeFileChunk:
			var chunk common.FileChunk
			if err := json.Unmarshal(msg.Payload, &chunk); err != nil {
				return fmt.Errorf("unmarshal chunk failed: %v", err)
			}

			// 更新状态
			status.FileID = chunk.FileID
			status.TotalChunks = chunk.TotalChunks
			status.CurrentChunk = chunk.ChunkIndex + 1

			// 保存分片
			chunks[chunk.ChunkIndex] = chunk.Data
			status.TotalSize += int64(len(chunk.Data))

			// 发送进度
			if progressChan != nil {
				progressChan <- status
			}

		case common.TypeFileComplete:
			// 检查是否收到所有分片
			if len(chunks) != status.TotalChunks {
				return fmt.Errorf("missing chunks: got %d, want %d", len(chunks), status.TotalChunks)
			}

			// 按顺序写入文件
			for i := 0; i < status.TotalChunks; i++ {
				data, ok := chunks[i]
				if !ok {
					return fmt.Errorf("missing chunk %d", i)
				}
				if _, err := file.Write(data); err != nil {
					return fmt.Errorf("write chunk failed: %v", err)
				}
			}

			status.Completed = true
			if progressChan != nil {
				progressChan <- status
			}
			return nil

		default:
			return fmt.Errorf("unexpected message type: %v", msg.Type)
		}
	}
}

// SendFile 发送文件
func (t *Tunnel) SendFile(fileID string, filePath string, progressChan chan<- FileTransferStatus) error {
	file, err := os.Open(filePath)
	if err != nil {
		return fmt.Errorf("open file failed: %v", err)
	}
	defer file.Close()

	fileInfo, err := file.Stat()
	if err != nil {
		return fmt.Errorf("get file info failed: %v", err)
	}

	if fileInfo.Size() > maxFileSize {
		return fmt.Errorf("file too large: %d bytes (max: %d)", fileInfo.Size(), maxFileSize)
	}

	totalChunks := (int(fileInfo.Size()) + common.ChunkSize - 1) / common.ChunkSize
	buffer := make([]byte, common.ChunkSize)
	status := FileTransferStatus{
		FileID:      fileID,
		TotalChunks: totalChunks,
	}

	for i := 0; i < totalChunks; i++ {
		n, err := file.Read(buffer)
		if err != nil && err != io.EOF {
			return fmt.Errorf("read file failed: %v", err)
		}

		chunk := &common.FileChunk{
			FileID:      fileID,
			ChunkIndex:  i,
			TotalChunks: totalChunks,
			ChunkSize:   n,
			Data:        buffer[:n],
		}

		chunkData, err := json.Marshal(chunk)
		if err != nil {
			return fmt.Errorf("marshal chunk failed: %v", err)
		}

		if err := t.SendMessage(&common.TunnelMessage{
			Type:    common.TypeFileChunk,
			Payload: chunkData,
		}); err != nil {
			return fmt.Errorf("send chunk failed: %v", err)
		}

		status.CurrentChunk = i + 1
		status.TotalSize += int64(n)
		if progressChan != nil {
			progressChan <- status
		}
	}

	// 发送完成消息
	if err := t.SendMessage(&common.TunnelMessage{
		Type:    common.TypeFileComplete,
		Payload: []byte(fileID),
	}); err != nil {
		return fmt.Errorf("send complete message failed: %v", err)
	}

	status.Completed = true
	if progressChan != nil {
		progressChan <- status
	}

	return nil
}
